# -*- coding: utf-8 -*-
# ==========================
# @Time    : 2023/12/6
# @Author  : zhoutengwei
# @File    : storage.py
# ==========================

import json
import time
import ray
import requests
from datetime import datetime
from ray.util.queue import Queue

from fastapi import FastAPI, HTTPException, BackgroundTasks
from pydantic import BaseModel
import asyncio
from modules import ASRPredictor, OCRPredictor
from modelscope.utils.config import Config
from loguru import logger

from storage import MinioStorage
from douyin_crawl import Douyin

# 配置日志信息
logger.add(
    "log_1.log",  # 日子文件名，后留一个占位符
    rotation='1 MB',     # 日志分割，可以根据时间也可以根据大小
    colorize=False,        # 是否染色
    level='DEBUG'         # 日志等级
)

app = FastAPI()
cfg = Config.from_file("configs/config.yaml")
default_app_cfg = cfg.applications[0]
deploy_cfg = default_app_cfg.deployments[1]
ray_actor_cfg = deploy_cfg.ray_actor_options
usr_cfg = deploy_cfg.user_config

s3_storage = MinioStorage(usr_cfg.s3_addr, usr_cfg.s3_access_key,
                          usr_cfg.s3_secret_key)

class TaskInput(BaseModel):
    task_id: str
    keywords: list[str]
    callback_url: str


def async_background_tasks(task_input: TaskInput):
    loop = asyncio.get_event_loop()
    loop.create_task(process_task_in_background(task_input))


def format_sources(sources: dict, task_id, cb_url):
    result = []
    for keyword, infos in sources.items():
        for info in infos:
            info.update({"keyword": keyword, "task_id": task_id, "cb_url": cb_url})
            result.append(info)

    return result


def process_task_in_background(task_input: TaskInput):
    # 这里可以添加你的任务处理逻辑
    # 例如，你可以使用task_id和keywords进行某些操作，并在完成后通过callback_url发送状态
    dou = Douyin()

    sources = {}
    callback_url = task_input.callback_url
    task_id = task_input.task_id
    for keyword in task_input.keywords:
        # 假设我们在这里处理每个关键词并获取结果
        items = dou.search_key(keyword, 0)
        sources[keyword] = []
        count = 0
        for item in items:
            #time.sleep(2)
            if item is None:
                continue
            sources[keyword].append(item)
            count += 1
            if len(sources[keyword]) > 2:
                break

    sources = format_sources(sources,task_id,callback_url)
    logger.info(sources)
    # with open("task_off.json", 'w') as f:
    #     json.dump(sources, f, ensure_ascii=False)

    ds = ray.data.from_items(sources, parallelism=1)
    asr_prediction = ds.map(ASRPredictor,
                            concurrency=1,
                            num_gpus=ray_actor_cfg.num_gpus,
                            fn_constructor_kwargs={'model_path': usr_cfg.model_path,
                                                   'chunk_size': usr_cfg.chunk_size,
                                                   'chunk_interval': usr_cfg.chunk_interval,
                                                   'sample_rate': usr_cfg.sample_rate,
                                                   'interval': usr_cfg.interval,
                                                   'min_silent_time': usr_cfg.min_silent_time,
                                                   'storage_obj': s3_storage})
    ocr_prediction = ds.map(OCRPredictor,
                            concurrency=1,
                            num_gpus=ray_actor_cfg.num_cpus,
                            fn_constructor_kwargs={'font_path': "resources/fonts/FZYTK.TTF",
                                                   'storage_obj': s3_storage})

    ocr_pred = ocr_prediction.take_all()
    asr_pred = asr_prediction.take_all()

    # 创建任务状态信息
    dataTypeToActions = {
        'base_info': sources,  # 基本信息
        'screen_shot': None,  # 截图并OCR
        'generate_txt': None,  # 转录
        'status': "finished"  # 状态
    }
    status_data = {
        "ts": datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
        "task_id": task_id,
        "status": "finished",
        "type": "task_off",
        "data": dataTypeToActions
    }
    # with open("task_off.json", 'w') as f:
    #     json.dump(status_data, f)
    # 使用callback_url发送结果
    send_status_to_callback(callback_url, status_data)


def send_status_to_callback(callback_url: str, result: dict):
    # 这里可以添加你的回调函数调用逻辑
    # 以下是一个简单的示例，使用aiohttp库发送POST请求
    print(result)
    # import aiohttp
    # async with aiohttp.ClientSession() as session:
    #     async with session.post(callback_url, json=result) as response:
    #         print(response.json())
    #         if response.status != 200:
    #             raise Exception(f"Failed to send status to callback: {response.status}")
    try:
        resp = requests.post(callback_url, json=result)
    except:
        logger.error("callback_url is not valid!")



@app.post("/process_task")
async def handle_process_task(task_input: TaskInput, background_tasks: BackgroundTasks):
    background_tasks.add_task(process_task_in_background, task_input)
    return {"message": f"Task {task_input.task_id} is being processed"}


if __name__ == "__main__":
    import uvicorn

    uvicorn.run(app, host="0.0.0.0", port=13036)
